#include "tinynet.h"

#include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <assert.h>

#ifdef WIN32
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <Winsock2.h>
#else
#endif

#define make(a) (a*)malloc(sizeof(a))

#define TN_OUTGOING_PACKET_BYTE_SIZE 2048
#define TN_INCOMING_PACKET_BYTE_SIZE 2048
#define TN_PROTOCOL_VERSION 'TNP1'

static BOOL sockets_started = FALSE;
static int last_error = 0;
static char* packet_buffer = 0;
static int packet_buffer_size = 0;

///************************************************************************/
/// Packet Queue
///************************************************************************/

typedef struct tn_packet_queue_element
{
	tn_packet* packet;
	struct tn_packet_queue_element* prev;
} tn_packet_queue_element;

struct tn_packet_queue
{
	tn_packet_queue_element* back;
	tn_packet_queue_element* served;
};

///************************************************************************/
/// Connection Internals
///************************************************************************/

struct tn_connection_internals
{
	uint32_t flags;
	tn_packet* incoming_packet;
	int incoming_packet_bytes_transmitted;
	struct tn_packet_queue outgoing_packet_queue;
	int outgoing_packet_bytes_transmitted;
};

typedef struct
{
  tn_connection connection;
  struct tn_connection_internals internals;
} tn_safe_connection_plus_internals;

///************************************************************************/
/// Helper Functions
///************************************************************************/

static tn_packet_queue_element* tn_pop_queue(struct tn_connection_internals* internals)
{
	tn_packet_queue_element* p = internals->outgoing_packet_queue.served;
	internals->outgoing_packet_queue.served = p->prev;
	if (p->prev == 0)
		internals->outgoing_packet_queue.back = 0;
	return p;
}

static void tn_push_queue(struct tn_connection_internals* internals, tn_packet_queue_element* qe)
{
	if (internals->outgoing_packet_queue.back)
		internals->outgoing_packet_queue.back->prev = qe;

	if (!internals->outgoing_packet_queue.served)
		internals->outgoing_packet_queue.served = qe;

	internals->outgoing_packet_queue.back = qe;
}

static tn_connection* tn_alloc_connection()
{
  tn_safe_connection_plus_internals* connection = 0;
  assert(&((tn_safe_connection_plus_internals*)0)->connection == (tn_safe_connection_plus_internals*)0);

  connection = make(tn_safe_connection_plus_internals);
	memset(connection, 0, sizeof(tn_safe_connection_plus_internals));

	connection->connection.internals = &connection->internals;
	return &connection->connection;
}

static int tn_determine_last_error(void)
{
#ifdef WIN32
	int last = WSAGetLastError();
	switch (last)
	{
	case WSANOTINITIALISED:
		return TNE_NOT_INITIALIZED;
	case WSAEMFILE:
		return TNE_NO_MORE_SOCKETS;
	case WSAEISCONN:
		return TNE_ALREADY_CONNECTED;
	case WSAENOTCONN:
		return TNE_NOT_CONNECTED;
	case WSAEWOULDBLOCK:
		return -1;
	default:
		return last * 100;
	}
#else
	return TNE_UNKNOWN_ERROR;
#endif
}

///************************************************************************/
/// API
///************************************************************************/

int tn_init(void)
{
  WORD wVersionRequested;
  WSADATA wsaData;
  int err = 0;

  if (sockets_started) return 1;
#ifdef WIN32
  wVersionRequested = MAKEWORD( 2, 2 );
  err = WSAStartup( wVersionRequested, &wsaData );
  if ( err != 0 )
	{
    /*
    int e = 0;
    if (e = gSockErrors(err)) return gPair(eSockError,e);
    return gPair(eSockError,err);*/
    return 0;
  }
#else
#endif

	atexit(tn_shutdown);

	last_error = TNE_OK;
	packet_buffer = (char*)malloc(64*1024);
	packet_buffer_size = 0;

  sockets_started = TRUE;

  return 1;
}

void tn_shutdown(void)
{
#ifdef WIN32
  if (sockets_started)
  {
    if (WSACleanup())
    {
      //int err = WSAGetLastError();
			WSAGetLastError();
      /*
      int e = 0;
      if (e = gSockErrors(err)) return gPair(eSockError,e);
      return gPair(eSockError,err);*/
      return;
    }
  } else return;
#endif

	free(packet_buffer);
	packet_buffer = 0;

  sockets_started = FALSE;
}

int tn_get_last_error(void)
{
	return last_error;
}

int tn_set_address(tn_address* dest, const char* host, uint16_t port)
{
	struct hostent* hbn = gethostbyname(host);
	struct in_addr* ina = (struct in_addr *)*hbn->h_addr_list;
	dest->ip = ina->s_addr;
	dest->port = port;

  return 1;
}

void tn_print_address(tn_address* address, char* buffer)
{
  int ipoct1 = address->ip&255;
  int ipoct2 = (address->ip>>8)&255;
  int ipoct3 = (address->ip>>16)&255;
  int ipoct4 = (address->ip>>24)&255;
  sprintf(buffer, "%i.%i.%i.%i:%i", ipoct1, ipoct2, ipoct3, ipoct4, address->port);
}

void tn_print_ip(uint32_t ip, char* buffer)
{
  int ipoct1 = ip&255;
  int ipoct2 = (ip>>8)&255;
  int ipoct3 = (ip>>16)&255;
  int ipoct4 = (ip>>24)&255;
  sprintf(buffer, "%i.%i.%i.%i", ipoct1, ipoct2, ipoct3, ipoct4);
}

tn_listener* tn_create_listener( tn_address* address )
{
	tn_listener* listener = 0;
	int status = 1;
	struct sockaddr_in sin;
	unsigned long iMode = 1;

	if (!address)
	{
		last_error = TNE_INCORRECT_PARAMETER;
		return 0;
	}

	listener = make(tn_listener);
	if (!listener)
	{
		last_error = TNE_OUT_OF_MEMORY;
		return 0;
	}

	if ((listener->socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0) 
	{
		last_error = tn_determine_last_error();
		free(listener);
		return 0;
	}

	setsockopt(listener->socket, SOL_SOCKET, SO_REUSEADDR, (char *)&status, sizeof(status)); /// reuse port
	// SO_EXCLUSIVEADDRUSE
	setsockopt(listener->socket, IPPROTO_TCP, TCP_NODELAY, (char *)&status, sizeof(status)); /// disable Nagle
	ioctlsocket(listener->socket, FIONBIO, &iMode); /// non-blocking

	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;
	sin.sin_addr.s_addr = INADDR_ANY;
	sin.sin_port = htons(address->port);

	if (bind(listener->socket, (struct sockaddr*)&sin, sizeof(sin)) != 0)
	{
		//int debug_error = tn_determine_last_error();
		last_error = TNE_CANNOT_REBIND_PORT;
		free(listener);
		return 0;
	}

	last_error = TNE_OK;

	listen(listener->socket, 128);

	listener->state = TNL_LISTENING;
	return listener;
}

void tn_destroy_listener( tn_listener* listener )
{
	if (!listener) return;
	shutdown(listener->socket, SD_BOTH);
	closesocket(listener->socket);
	free(listener);
}


enum tn_internal_flags
{
	/// Connection flags
	TNI_CONNECTION_TO_CLIENT = 0,
	TNI_CONNECTION_TO_SERVER,

	/// Packet flags
	TNI_OWNS_DATA = 0,
};

tn_connection* tn_listen( tn_listener* listener )
{
	tn_connection* connection;
	int addrlen = sizeof(struct sockaddr_in);
	intptr_t socket = -1;
	int temp = htonl(TN_PROTOCOL_VERSION);
	struct sockaddr_in pin;

	if (!listener)
	{
		last_error = TNE_INCORRECT_PARAMETER;
		return 0;
	}
	if (listener->state != TNL_LISTENING)
	{
		last_error = TNE_LISTENER_NOT_ACTIVE;
		return 0;
	}

	if ((socket = (intptr_t)accept(listener->socket, (struct sockaddr*)&pin, &addrlen)) < 0)
	{
		int lerr = tn_determine_last_error();
		if (lerr == -1) /// nothing yet
		{ 
			last_error = TNE_OK;
			return 0;
		} else
		{
			last_error = lerr;
			return 0;
		}
	}

	last_error = TNE_OK;

	connection = tn_alloc_connection();
	if (!connection) 
	{
		last_error = TNE_OUT_OF_MEMORY;
		return 0;
	}

	connection->state = TNC_ESTABLISHING;
	connection->remote_ip = pin.sin_addr.s_addr;
	connection->socket = socket;
	connection->internals->flags = 1<<TNI_CONNECTION_TO_CLIENT;

	tn_send_data(connection, &temp, 4);

	return connection;
}

static void tn_disconnect(tn_connection* connection)
{
	char temp[4];
	time_t disconnect_time = 0;
	shutdown(connection->socket, SD_BOTH);

	disconnect_time = time(0);
	while (recv(connection->socket, temp, 4, 0) > 0 && time(0)-disconnect_time < 100) ;
	closesocket(connection->socket);

	connection->state = TNC_DISCONNECTED;
}

void tn_destroy_connection( tn_connection* connection )
{
	if (!connection) return;

	while (connection->internals->outgoing_packet_queue.served)
	{
		tn_packet_queue_element* p = tn_pop_queue(connection->internals);
		tn_destroy_packet(p->packet);
		free(p);
	}

	if (connection->internals->incoming_packet)
		tn_destroy_packet(connection->internals->incoming_packet);

	tn_disconnect(connection);

	free(connection);
}

tn_connection* tn_connect_to( tn_address* address )
{
	tn_connection* connection = 0;
	struct sockaddr_in pin;
	int status = 1;
	unsigned long iMode = 1;
	int temp = htonl(TN_PROTOCOL_VERSION);
	//int result = 0, optres = -1, optlen = sizeof(int);

	if (!address)
	{
		last_error = TNE_INCORRECT_PARAMETER;
		return 0;
	}

	connection = tn_alloc_connection();
	if (!connection)
	{
		last_error = TNE_OUT_OF_MEMORY;
		return 0;
	}

	if ((connection->socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
	{
		last_error = tn_determine_last_error();
		free(connection);
		return 0;
	}


	if (setsockopt(connection->socket, IPPROTO_TCP, TCP_NODELAY, (char*)&status, sizeof(status))) /// disable Nagle
	{
		last_error = tn_determine_last_error();
		free(connection);
		return 0;
	}

	if (ioctlsocket(connection->socket, FIONBIO, &iMode)) /// non-blocking
	{
		last_error = tn_determine_last_error();
		free(connection);
		return 0;
	}

	memset(&pin, 0, sizeof(pin));
	pin.sin_family = AF_INET;
	pin.sin_addr.s_addr = address->ip;
	pin.sin_port = htons(address->port);

	connection->remote_ip = address->ip;
	connection->internals->flags = 1<<TNI_CONNECTION_TO_SERVER;
	connection->state = TNC_CONNECTING;

	if (connect(connection->socket, (struct sockaddr *)&pin, sizeof(pin)) < 0)
	{
		int lerr = tn_determine_last_error();
		if (lerr != -1)
		{
			last_error = lerr;
			free(connection);
			return 0;
		} else
		{
			last_error = TNE_OK; /// non-blocking connection underway
			return connection;
		}
	}

	/// non-blocking connection established immediately

	last_error = TNE_OK;

	connection->state = TNC_ESTABLISHING;

	tn_send_data(connection, &temp, 4);

	return connection;
}

enum tn_internal_packet_flags
{
	TNIP_IN_QUEUE,
	TNIP_SENDING,
	TNIP_RECEIVING,
	TNIP_OWNS_DATA,
};

static void tn_work_connecting(tn_connection* connection)
{
	fd_set success_set, failure_set;
	struct timeval timeout;
	int how_many = 0;
	int temp = htonl(TN_PROTOCOL_VERSION);

	timeout.tv_sec = timeout.tv_usec = 0;

	FD_ZERO(&success_set);
	FD_ZERO(&failure_set);
	FD_SET(connection->socket, &success_set);
	FD_SET(connection->socket, &failure_set);

	how_many = select(FD_SETSIZE, (fd_set*)0, &success_set, &failure_set, &timeout);
	if (how_many == 1) /// connected or disconnected
	{
		last_error = TNE_OK;

		if (FD_ISSET(connection->socket, &success_set))
		{
			connection->state = TNC_ESTABLISHING;
			tn_send_data(connection, &temp, 4);
		}
		else if (FD_ISSET(connection->socket, &failure_set))
			connection->state = TNC_COULD_NOT_CONNECT;
	} else if (how_many == 0) /// not connected yet
	{
		last_error = TNE_OK;
	} else /// uuu, error
	{
		last_error = tn_determine_last_error();
	}
}

static tn_packet* tn_work_working( tn_connection* connection )
{
	/// Recv
	if (!connection->internals->incoming_packet) /// try to receive packet size
	{
		int16_t peek_size = -1;
		int recv_bytes = recv(connection->socket, (char*)&peek_size, sizeof(peek_size), MSG_PEEK);
		if (recv_bytes < 0) /// error or would-block
		{
			int lerr = tn_determine_last_error();
			if (lerr == -1) /// no data on queue
			{
				
			} else
			{
				last_error = lerr;
				return 0;
			}
		} 
		else if (recv_bytes == 0) /// no bytes, other side disconnected
		{
			connection->state = TNC_DISCONNECTED;
		} 
		else if (recv_bytes == sizeof(uint16_t)) /// got packet size
		{
			tn_packet* packet = make(tn_packet);

			peek_size = ntohs(peek_size);
			packet->size = peek_size;
			packet->internal_flags = (1<<TNIP_RECEIVING)|(1<<TNIP_OWNS_DATA);
			packet->data = (char*)malloc(peek_size);

			connection->internals->incoming_packet = packet;
			connection->internals->incoming_packet_bytes_transmitted = 0;

			recv(connection->socket, (char*)&peek_size, sizeof(peek_size), 0); /// will succeed
		} 
		else /// no packet size yet
		{
			/// nothing here
		}
	}

	if (connection->internals->incoming_packet) /// receive bytes
	{
		struct tn_connection_internals* internals = connection->internals;
		tn_packet* incoming_packet = internals->incoming_packet;
		int packet_bytes_transmitted = internals->incoming_packet_bytes_transmitted;

		int recv_bytes = 0;

		if (incoming_packet->size > 0)
		{
      /// either TN_INCOMING_PACKET_BYTE_SIZE bytes, or the rest of the packet
			int bytes_to_get = min(incoming_packet->size-packet_bytes_transmitted, TN_INCOMING_PACKET_BYTE_SIZE);

			recv_bytes = recv(connection->socket, incoming_packet->data + (ptrdiff_t)packet_bytes_transmitted, bytes_to_get, 0);

			if (recv_bytes < 0)
			{
				int lerr = tn_determine_last_error();
				if (lerr == -1) /// no data on queue
				{
				} else
				{
					last_error = lerr;
					return 0;
				}
			} else if (recv_bytes == 0) /// graceful disconnect
			{
				connection->state = TNC_DISCONNECTED;
				return 0;
			}
		}
		
		internals->incoming_packet_bytes_transmitted += recv_bytes;
		if (packet_bytes_transmitted + recv_bytes >= incoming_packet->size) /// packet done!
		{
			internals->incoming_packet = 0;
			internals->incoming_packet_bytes_transmitted = 0;

			incoming_packet->internal_flags &= ~(1<<TNIP_RECEIVING);

			return incoming_packet;
		}
	}

	/// Send
	if (connection->internals->outgoing_packet_queue.served)
	{
		struct tn_connection_internals* internals = connection->internals;
		tn_packet* outgoing_packet = internals->outgoing_packet_queue.served->packet;
		int packet_bytes_transmitted = internals->outgoing_packet_bytes_transmitted;

		int sent_bytes = -1;
		/// either TN_OUTGOING_PACKET_BYTE_SIZE bytes, or the rest of the packet
		int bytes_to_transmit = min(outgoing_packet->size-packet_bytes_transmitted, TN_OUTGOING_PACKET_BYTE_SIZE);

		outgoing_packet->internal_flags &= ~(1<<TNIP_IN_QUEUE);
		outgoing_packet->internal_flags |= 1<<TNIP_SENDING;

		sent_bytes = send(connection->socket, outgoing_packet->data + packet_bytes_transmitted, bytes_to_transmit, 0); /// send data

		if (sent_bytes < 0)
		{
			int lerr = tn_determine_last_error();
			if (lerr == -1) /// should not happen
			{
				last_error = TNE_OK;
				assert(0 && "Should not happen!");
				return 0;
			} else
			{
				last_error = lerr;
				return 0;
			}
		}

		/// add transmitted size
		internals->outgoing_packet_bytes_transmitted += sent_bytes;
		if (packet_bytes_transmitted + sent_bytes >= outgoing_packet->size) /// packet transmitted!
		{
			tn_packet_queue_element* pe = tn_pop_queue(internals);
			tn_destroy_packet(outgoing_packet);
			free(pe);

			internals->outgoing_packet_bytes_transmitted = 0;
		}
	}

	last_error = TNE_OK;
	return 0;
}

tn_packet* tn_work_connection( tn_connection* connection )
{
	tn_packet* return_packet = 0;

	if (!connection)
	{
		last_error = TNE_INCORRECT_PARAMETER;
		return 0;
	}
	if (connection->state <= 0)
	{
		last_error = TNE_NOT_CREATED;
		return 0;
	}

	switch (connection->state)
	{
	case TNC_CONNECTING:
		tn_work_connecting(connection);
		return 0;
	case TNC_ESTABLISHING:
		return_packet = tn_work_working(connection);
		if (return_packet)
		{
			if (return_packet->size == 4 && ntohl(*((uint32_t*)return_packet->data)) == TN_PROTOCOL_VERSION)
			{
				connection->state = TNC_WORKING;
				last_error = TNE_OK;
			} else
			{
				tn_disconnect(connection);
				last_error = TNE_WRONG_NETWORK_PROTOCOL;
			}
			tn_destroy_packet(return_packet);
		}
		return 0;
	case TNC_WORKING:
		return tn_work_working(connection);
	default:
		last_error = TNE_INCORRECT_PARAMETER;
		return 0;
	}
}

void tn_destroy_packet(tn_packet* packet)
{
	if (!packet) return;

	if (packet->internal_flags & (1<<TNIP_OWNS_DATA))
	{
		free(packet->data);
	}
	free(packet);
}

char* tn_release_packet_data( tn_packet* packet )
{
	char* data = 0;
	if (!packet)
	{
		last_error = TNE_INCORRECT_PARAMETER;
		return 0;
	}

	data = packet->data;
	packet->internal_flags &= ~(1<<TNIP_OWNS_DATA);

	return data;
}

int tn_send_data( tn_connection* connection, void* data, uint16_t size )
{
	tn_packet* pack = 0;
	tn_packet_queue_element* qe = 0;
	char* cdata = (char*)data;
	if (!connection || !data || !size)
	{
		last_error = TNE_INCORRECT_PARAMETER;
		return -1;
	}

	/// Create packet
	pack = make(tn_packet);
	pack->size = size + 2;
	pack->data = malloc(size + 2);
	memcpy(pack->data+2, cdata, size);
	*((uint16_t*)pack->data) = htons(size);
	pack->internal_flags = 1<<TNIP_OWNS_DATA;

	qe = make(tn_packet_queue_element);
	qe->packet = pack;
	qe->prev = 0;

	tn_push_queue(connection->internals, qe);

	return TNE_OK;
}

int tn_send_packet( tn_connection* connection, tn_packet* packet )
{
  return tn_send_data(connection, packet->data, packet->size);
}

int tn_send_string( tn_connection* connection, const char* string )
{
	return tn_send_data(connection, (void*)string, (uint16_t)(strlen(string)+1));
}